#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np
import os

def range_custom(start, end, step, output_dtype=np.float32):
    """自定义 range 实现, 生成 [start, end) 间隔 step 的等差数列.

    支持:
        - Python 标量 (int/float)
        - numpy 标量/张量 (0 维或 1 维, 如 shape=() 或 (1,))
    最终都会转换成 float32 标量参与计算.
    """

    # 统一转为 float32 标量
    start_f32 = np.float32(start)
    end_f32 = np.float32(end)
    step_f32 = np.float32(step)

    # step 不能为 0
    if step_f32 == 0:
        raise ValueError("step 不能为 0")

    # 计算输出长度: ceil(abs(end - start) / abs(step))
    diff = np.abs(end_f32 - start_f32)
    step_abs = np.abs(step_f32)

    if diff == 0:
        num_elements = 0
    else:
        num_elements = int(np.ceil(diff / step_abs))

    # 调试信息：显示关键参数和计算结果
    print(
        f"[gen_data] Range parameters: "
        f"start={start_f32}, end={end_f32}, step={step_f32}, "
        f"diff={diff}, step_abs={step_abs}, num_elements={num_elements}"
    )

    if num_elements <= 0:
        # 返回空 ND 数组
        return np.array([], dtype=output_dtype)

    # 生成序列: start + i * step, i = 0..num_elements-1
    seq = start_f32 + np.arange(num_elements, dtype=np.float32) * step_f32
    return seq.astype(output_dtype)

def get_file_size_bytes(file_path):
    """获取文件大小"""
    if os.path.exists(file_path):
        return os.path.getsize(file_path)
    return 0

def gen_golden_data_range():
    np.random.seed()

    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    # ==================== 固定数据类型配置 ====================
    input_dtype_1 = np.float32    
    # input_dtype_1 = np.float32
    # input_dtype_1 = np.int16
    # input_dtype_1 = np.int32
    input_dtype_2 = np.float32    
    # input_dtype_2 = np.float32
    # input_dtype_2 = np.int16
    # input_dtype_2 = np.int32
    output_dtype = np.float32   
    ns_dtype = np.float32         
    # =======================基础参数配置======================
    min_step = -100.0                    # step的最小值
    max_step = 100.0                     # step的最大值
    start_range = (-20, 20)              # start的取值范围
    end_offset_range = (-1000, 10000)    # end与start的差值范围
    # ========================================================
    # 生成 start
    starts = np.random.uniform(*start_range, size=1).astype(input_dtype_1)
    # 生成 end = start + offset
    end_offset = np.random.uniform(*end_offset_range, size=1).astype(input_dtype_2)
    ends = input_dtype_2(starts + end_offset)
    # 生成 step, 保证不为 0 且绝对值不过小
    min_abs_step = 1e-3
    while True:
        ns = np.random.uniform(min_step, max_step, size=1).astype(ns_dtype)
        if np.abs(ns[0]) >= min_abs_step:
            break

    golden = range_custom(starts[0], ends[0], ns[0], output_dtype=output_dtype)
    
    starts.tofile("./input/starts.bin")
    ends.tofile("./input/ends.bin")
    ns.tofile("./input/ns.bin")
    golden.tofile("./output/golden.bin")
    
    # 打印调试信息
    print("调试信息:")
    print(f"start值: {starts[0]} (类型: {starts.dtype})")
    print(f"end值: {ends[0]} (类型: {ends.dtype})")
    print(f"step值: {ns[0]} (类型: {ns.dtype})")
    print(f"生成的数据长度: {len(golden)}")
    print(f"生成的数据类型: {golden.dtype}")
    print("\n文件大小 (字节):")
    start_size = get_file_size_bytes("./input/starts.bin")
    print(f"starts.bin: {start_size} bytes")
    
    end_size = get_file_size_bytes("./input/ends.bin")
    print(f"ends.bin: {end_size} bytes")
    
    ns_size = get_file_size_bytes("./input/ns.bin")
    print(f"ns.bin: {ns_size} bytes")
    
    golden_size = get_file_size_bytes("./output/golden.bin")
    print(f"golden.bin: {golden_size} bytes")

if __name__ == "__main__":
    gen_golden_data_range()